import os
import sys
import cv2
import shutil
import pickle
import argparse
import numpy as np

sys.path.append(
    os.path.realpath(os.path.join(os.path.dirname(__file__), '..')))
from perception.common.video import clip_video_to_frames, get_video_extra_info
from perception.utterance.eval import UtteranceEncoder
from interaction.common.data import XiaoduHiDataset, \
    SalutationClsDataset, check_passive_interaction
from interaction.common.data_via_decord import XiaoduHiDecord


def parse_args():
    parser = argparse.ArgumentParser(
        description='Split and prepare annotions for XiaoduHi dataset.')
    parser.add_argument('--data_version', '-dv', type=str, default='ds',
                        help='The version of the dataset.')
    parser.add_argument('--output_dir', '-o', type=str, default='data',
                        help='The output directory to store generated '
                        'training and testing dataset files.')
    parser.add_argument(
        '--anno_dir', '-ad', type=str, default='data/annos',
        help='Directory of V2 annotations.')
    parser.add_argument(
        '--video_tracking_dir', '-vd', type=str, default='data/clips',
        help='Directory of preprocessed video obj-tracking results.')

    data = parser.add_argument_group('data')
    data.add_argument(
        '--wae_dir', '-wd', type=str, default='data/raw_wae',
        help='Directory of multimodal action embeddings, generated by '
        'scripts/collect_v2_act_emb.py.')

    decord = parser.add_argument_group('decord')
    decord.add_argument(
        '--train_pkl', type=str, default='data/train.pkl',
        help='Path to v2 training dataset pkl.')
    decord.add_argument(
        '--test_pkl', type=str, default='data/test.pkl',
        help='Path to v2 test dataset pkl.')

    yolov4_group = parser.add_argument_group('YOLOv4')
    yolov4_group.add_argument(
        '--yolov4_model_dir', type=str,
        default='tools/yolov4_paddle/inference_model',
        help='Directory of YOLOv4 model.')

    return parser.parse_args()


def prepare_datasets(args):
    dataset = XiaoduHiDataset(
        args.video_tracking_dir, args.anno_dir, args.wae_dir)
    dataset.build_dataset(args.output_dir)


def prepare_dataset_decord(args):
    dataset = XiaoduHiDecord(
        args.video_tracking_dir, args.train_pkl, args.test_pkl, args.anno_dir)
    dataset.build_dataset(args.output_dir, workers=32)


def prepare_salutation_dataset(args):
    dataset = SalutationClsDataset(
        args.video_tracking_dir, args.anno_dir, args.yolov4_model_dir)
    dataset.build_dataset(args.output_dir)


if __name__ == '__main__':
    if len(sys.argv) == 1:
        sys.argv.append('-h')
    args = parse_args()

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    if args.data_version == 'ds':
        prepare_datasets(args)
    elif args.data_version == 'ds_decord':
        prepare_dataset_decord(args)
    elif args.data_version == 'salutation_v2':
        prepare_salutation_dataset(args)
    else:
        raise ValueError('Invalid data version')
